from typing import Dict, Any
from langgraph.graph import StateGraph, END
from src.core.state import AgentState
from src.agents.router import AgentRouter
from src.agents.jira_agent import JiraAgent
from src.agents.conference_agent import ConferenceAgent
import logging

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class GraphManager:
    """
    Manages the LangGraph workflow for the agent system.
    """
    
    def __init__(self):
         # 初始化GraphManager，创建所需的代理和图构建器
        logger.info("Initializing GraphManager")
        self.router = AgentRouter() # 路由代理，负责请求分发
        self.jira_agent = JiraAgent()  # Jira代理，处理Jira相关操作
        self.conference_agent = ConferenceAgent()
        self.builder = StateGraph(AgentState) # 基于AgentState的状态图构建器
        self._setup_graph() # 设置图结构
    
    def _setup_graph(self):
        """
        Set up the graph nodes and edges.
        """
        logger.info("Setting up graph nodes and edges")
        
        # 添加节点
        self.builder.add_node("router", self._route_request) # 路由节点
        self.builder.add_node("jira_agent", self._process_jira_request) # Jira处理节点
        self.builder.add_node("conference_agent", self._process_conference_request)
        
        # 添加条件边，根据_route_decision函数输出决定路由
        self.builder.add_conditional_edges(
            "router",
            self._route_decision,
            {
                "jira_agent": "jira_agent",
                "conference_agent": "conference_agent",
                "end": END
            }
        )
        
        # 添加从Agent到工作流结束的线性边
        self.builder.add_edge("jira_agent", END)
        self.builder.add_edge("conference_agent", END)
        
        # 设置入口点，所有工作流都从路由器开始
        self.builder.set_entry_point("router")
        logger.info("Graph setup completed")
    
    def _route_request(self, state: AgentState) -> Dict[str, Any]:
        """
        Route the request to the appropriate agent.
        """
        logger.info("Routing request")
        return self.router.process(state)
    
    def _process_jira_request(self, state: AgentState) -> Dict[str, Any]:
        """
        Process a Jira request.
        """
        logger.info("Processing Jira request")
        return self.jira_agent.process(state)
    
    def _process_conference_request(self, state: AgentState) -> Dict[str, Any]:
        """
        Process a conference request.
        """
        logger.info("Processing Conference request")
        return self.conference_agent.process(state)
    
    def _route_decision(self, state: AgentState) -> str:
        """
        Make routing decision based on state.
        """
        # 访问最后一条消息以确定路由
        messages = state.get("messages", [])
        logger.info(f"Making routing decision based on state. Messages count: {len(messages)}")
        
        if not messages:
            logger.info("No messages found, routing to end")
            return "end"
            
        # 目前根据current_agent字段进行路由
        current_agent = state.get("current_agent", "")
        logger.info(f"Current agent: {current_agent}")
        
        if current_agent == "jira":
            # 路由到Jira代理
            logger.info("Routing to jira_agent")
            return "jira_agent"
        elif current_agent == "conference":
            logger.info("Routing to conference_agent")
            return "conference_agent"
        else:
            logger.info("Routing to end")
            return "end"
    
    def compile(self):
        """
        Compile and return the graph.
        """
        logger.info("Compiling graph")
        return self.builder.compile()